import os
import numpy as np
from glob import glob
import os.path as osp
from PyVGrid.vgrid import VGrid
from tqdm import tqdm
import xml.etree.ElementTree as ET
import shutil
from util.vgrid_util import VGridParam

def makesure_dir(_path):
    if not osp.exists(_path):
        os.makedirs(_path)

def getSizeFormIQ(iq_path):
    tree = ET.parse(iq_path)
    root = tree.getroot()
    ImageParam = root.find('ImageParam')
    height = ImageParam.find('Height').text
    width = ImageParam.find('Width').text

    PyramidParam = root.find('PyramidParam')
    PContrast = PyramidParam.find('PContrast').text
    DContrast = PyramidParam.find('DContrast').text
    BContrast = PyramidParam.find('BContrast').text
    BrightNess = PyramidParam.find('BrightNess').text

    WLGParam = root.find('WLGParam')
    Gamma = WLGParam.find('Gamma').text
    return int(height), int(width), float(PContrast), float(DContrast), float(BContrast), float(BrightNess), float(Gamma)


if __name__ == '__main__':
    base_path_list = [
        r"F:\不同部位数据集\腹部\腹部-娄底",
        r"F:\不同部位数据集\腹部\腹部-裕和",
        r"F:\不同部位数据集\腹部\腹部-未知"
    ]
    save_path_list = [
        "F:\不同部位数据集对\腹部\腹部-娄底\VGRID_B54_NOSNR",
        "F:\不同部位数据集对\腹部\腹部-裕和\VGRID_B54_NOSNR",
        "F:\不同部位数据集对\腹部\腹部-未知\VGRID_B54_NOSNR",
    ]
    for base_path, save_path in zip(base_path_list, save_path_list):
        base_path = r"G:\不同部位数据集\胸部\胸部-裕和"
        save_path = r"G:\不同部位数据集对\胸部-裕和\VGRID_Gamma(原始效果fromIQ)"
        makesure_dir(save_path)
        name_list = [path.split(".raw")[0] for path in glob(osp.join(base_path, "**", "*.raw"), recursive=True)]
        tqdm_bar = tqdm(name_list, total=len(name_list))
        vg = VGrid()
        for name in tqdm_bar:
            base_name = osp.basename(name)
            iq_path = f"{name}.iq"
            raw_path = f"{name}.raw"
            if not osp.exists(raw_path) or not osp.exists(iq_path):
                print(f"have not paired raw or iq file, {name}")
                continue
            height, width, p, d, b, BT, gamma = getSizeFormIQ(iq_path)
            vgParam = VGridParam()
            vgParam.PContrast, vgParam.DContrast, vgParam.BContrast, vgParam.brightness = 3, 3.5, 3.5, 0.54
            vgParam.SNRFlag, vgParam.lpSubFlag, vgParam.histFlag = False, False, False
            vgParam.acqMode, vgParam.acqChildmode, vgParam.bodyPart = 1, 0, 5

            imgs = np.fromfile(raw_path, dtype=np.uint16).reshape((-1, height, width))
            frame_count = len(imgs)
            vg.Initialization(width, height, 16)

            result_imgs = []
            for frame_id, img in enumerate(imgs):
                img_out = np.zeros_like(img)
                # vg.Process(img, img_out, p, d, b, 1.0, BT, 0,
                #            True, False, False, width, height, False,
                #            1, 0, 1,
                #            20, 15, 0, 0, frame_id+1, 0, False
                #            )
                vg.Process(img, img_out, vgParam.PContrast, vgParam.DContrast, vgParam.BContrast, 1.0, vgParam.brightness, 0,
                           vgParam.SNRFlag, vgParam.lpSubFlag, vgParam.histFlag, width, height, False,
                           vgParam.acqMode, vgParam.acqChildmode, vgParam.bodyPart,
                           20, 15, 0, 0, frame_id+1, 0, False
                           )
                # img_out = np.clip(img_out, 0, 65535)
                # img_out = np.power(img_out.astype(np.float32) / 65535, gamma)           # gamma step 1
                # img_out = (img_out * 65535).astype("uint16")                            # gamma step 2
                result_imgs.append(img_out)
                tqdm_bar.set_postfix_str(f"{frame_id+1}/{frame_count}")

            result_imgs = np.array(result_imgs, np.uint16)
            # result_save_path = osp.join(save_path, name)
            # makesure_dir(result_save_path)
            # result_imgs.tofile(osp.join(result_save_path, name+".raw"))       # 使用文件夹包括raw文件
            # shutil.copy(iq_path, result_save_path)
            result_imgs.tofile(osp.join(save_path, base_name+".raw"))
            shutil.copy(iq_path, save_path)


            # for BT in [0.3, 0.4, 0.5, 0.6, 0.7]:
            #     result_imgs = []
            #     for frame_id, img in enumerate(imgs):
            #         img_out = np.zeros_like(img)
            #         vg.Process(img, img_out, p, d, b, 1.0, BT, 0,
            #                    True, True, False, width, height, False,
            #                    1, 0, 1,
            #                    20, 15, 0, 0, frame_id + 1, 0, False
            #                    )
            #         img_out = np.clip(img_out, 0, 65535)
            #         result_imgs.append(img_out)
            #         tqdm_bar.set_postfix_str(f"{frame_id + 1}/{frame_count}")
            #     result_save_path = osp.join(save_path, name)
            #     makesure_dir(result_save_path)
            #     result_imgs = np.array(result_imgs, np.uint16)
            #     result_imgs.tofile(osp.join(result_save_path, f"{name}_b{int(BT*10)}.raw"))
            #     shutil.copy(iq_path, result_save_path)







